import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
import argparse


import torch
import random
import numpy as np
import os


def get_params():

    parser = argparse.ArgumentParser()

    parser.add_argument('--seed', type=int, default=369, help='Random seed.')

    # parameters for gcl learning
    parser.add_argument('--epochs_gcl', type=int, default=100, help='Number of epochs to train gcl model.')
    parser.add_argument('--gcl_lr', type=float, default=0.01, help='Initial learning rate for gcl learning.')
    parser.add_argument('--input_dim_encoder', type=int, default=64, help='Number of input units.')
    parser.add_argument('--hidden_dim_encoder', type=int, default=32, help='Number of hidden units.')
    parser.add_argument('--num_layers_encoder', type=int, default=2, help='Number of layers for gnn encoder.')

    # for solve on graphs
    parser.add_argument('--inner_num_to_solve', type=int, default=100,
                        help='inner time to solve')
    parser.add_argument('--out_num_to_solve', type=int, default=10000,
                        help='outsied time to solve')
    parser.add_argument('--Tolerance_time', type=float, default=300.0,
                        help='Tolerance_time to solve')

    # parameters for dataset
    # data infor
    parser.add_argument('--train_ratio', type=float, default=0.8,
                        help='train_ratio.')
    parser.add_argument('--data_name', type=str,
                        help='may be crg_gnp_random_graph; rpt_rt_tree_graph; rc_bg_graph')
    parser.add_argument('--graph_name', type=str,default='crg_gnp_0.2',
                        help='may be crg_gnp_p, p=0.2~0.9; rpt_rt; rc_bg')
    parser.add_argument('--label_list', type=list, default=['crg', 'gnp'],
                        help="may be ['crg', 'gnp'], ['rpt', 'rt'], ['rc', 'bg']")
    # parameters for sampler training
    # parameters for sampler
    parser.add_argument('--time_step', type=int, default=3,
                        help='length of a trace.')
    parser.add_argument('--with_feature', type=bool, default=True,
                        help='True or false')
    parser.add_argument('--use_feature', type=str, default='use_feature',
                        help='use_feature,no_feature')
    parser.add_argument('--use_gnn', type=bool, default=False,
                        help='True or False')

    parser.add_argument('--hidden_dim', type=int, default=64, help='Number of hidden units for sampler.')
    parser.add_argument('--gamma', type=int, default=0.98, help='discount factor for compute the reward.')
    parser.add_argument('--reward_fun_number', type=int,default=1,
                        help='number to train the ls sampler:[1~8]')

    parser.add_argument('--num_episodes', type=int, default=1000, help='episodes for training.')
    parser.add_argument('--early_stop', type=int, default=50,
                        help='Initial learning rate.')
    parser.add_argument('--feature_dim', type=int, default=64,
                        help='feature dimension.')

    # sampling  feature randomly
    parser.add_argument('--l', type=float, default=-1.,
                        help='low bound of uniform distribution.')
    parser.add_argument('--u', type=float, default=1.,
                        help='up bound of uniform distribution.')


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument('--device', type=str, default=device,
                        help='device to use.')


    # for PPO sampler
    parser.add_argument('--ppo_lmbda', type=float, default=0.95,
                        help='for gae.')

    # parser.add_argument('--ppo_alpha', type=float, default=0.5,
    #                     help='for line search.')
    parser.add_argument('--ppo_inner_epochs', type=int, default=3,
                        help='for line search.')
    parser.add_argument('--ppo_eps', type=float, default=0.2,
                        help='for line search.')
    parser.add_argument('--rau', type=float, default=0.95,
                        help='ant.')
    parser.add_argument('--add_iter', type=int, default=1,
                        help='ant.')
    parser.add_argument('--reduce_iter', type=int, default=1,
                        help='ant.')
    parser.add_argument('--low_pher', type=float, default=0.,
                        help='ant.')
    parser.add_argument('--high_pher', type=float, default=1000.,
                        help='ant.')

    parser.add_argument('--critic_lr', type=float, default=0.01,
                        help='ppo.')
    parser.add_argument('--actor_lr', type=float, default=0.01,
                        help='ppo.')
    parser.add_argument('--env_lr', type=float, default=0.01,
                        help='ant.')
    parser.add_argument('--exp_name', type=str, default='salmas',
                        help='type of experiment.')
    parser.add_argument('--step', type=str, default='single_train',
                        help='type of learning or solving.')
    parser.add_argument('--dataset', type=str, default='salmas_data_1',
                        help='total dataset.')
    parser.add_argument('--max_norm', type=float, default=60.,
                        help='for grad clip.')



    args = parser.parse_args()

    return args


def set_seed(seed):
    """Set seed"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.enable = True
    os.environ["PYTHONHASHSEED"] = str(seed)


args = get_params()
set_seed(args.seed)

